package st;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import lm.models.BigramModel;

import st.utils.LexicalProb;

public class Viterbi {

	private BigramModel bigramProb;

	private List<String> tagSet;
	private List<String> inputWords;

	private static final String START = "#START#";

	public Viterbi(BigramModel bigramProb, List<String> tagSet,
			List<String> inputWords) {
		this.bigramProb = bigramProb;
		this.tagSet = tagSet;
		this.inputWords = inputWords;

	}

	public List<String> runViterbi() {
		List<String> tags = new ArrayList<String>();

		// lexical probabilities initialization

		int c = tagSet.size();
		int n = inputWords.size();
		int bptr[][] = new int[c][n];
		double score[][] = new double[c][n];

		// Initialization
		/*
		 * For i = 1 to c do SCORE(i,1) = P(ti|O) P(w1|ti) BPTR(i,1) = 0
		 */

		System.out.println("-------Initialization-------");
		for (int i = 0; i < c; i++) {
			if (tagSet.get(i).equalsIgnoreCase("<s>")) {
				score[i][0] = Math.log(((double) 2 / (double) c))
						+ Math.log(LexicalProb.getLexicalProb(
								inputWords.get(0), tagSet.get(i)));
			} else {
				score[i][0] = Math.log(((double) 1 / (double) c))
						+ Math.log(LexicalProb.getLexicalProb(
								inputWords.get(0), tagSet.get(i)));
			}

			System.out.println("Init = " + score[i][0]);
			bptr[i][0] = 0;

		}

		System.out.println("-------Build Score and BPTR-------");

		int maxIndex = 0;
		// Iteration
		for (int t = 1; t < n; t++) {

			System.out.println(t);
			for (int i = 0; i < c; i++) {
				score[i][t] = 0;

				// get max score and index
				for (int j = 0; j < c; j++) {
					double temp = score[j][t - 1]
							+ Math.log(bigramProb
									.getSmoothedBigramProbability(tagSet.get(j)
											+ " " + tagSet.get(i)))
							+ Math.log(LexicalProb.getLexicalProb(inputWords
									.get(t), tagSet.get(i)));
//					System.out.println("value for t= "
//							+ t
//							+ " i="
//							+ i
//							+ " is temp = "
//							+ temp
//							+ " score = "
//							+ score[j][t - 1]
//							+ " bigram="
//							+ bigramProb.getSmoothedBigramProbability(tagSet
//									.get(j)
//									+ " " + tagSet.get(i))
//							+ " lexical= "
//							+ LexicalProb.getLexicalProb(inputWords.get(t),
//									tagSet.get(i)));
					if (score[j][t - 1] == 0) {
//						System.out.println("Zero!!!");

					}

					if (score[i][t] == 0 || score[i][t] < temp) {
						score[i][t] = temp;
						maxIndex = j;
					}

				}
//				System.out.println("MAX = " + score[i][t]);
				// System.out.println("maxIndex = " + maxIndex);
				bptr[i][t] = maxIndex;
			}
		}

		System.out.println("-------Identify sequence-------");
		// Identify sequence
		double temp = 0;
		int indexForMax = 0;
		for (int i = 0; i < c; i++) {
			if (temp < score[i][n - 1]) {
				temp = score[i][n - 1];
				indexForMax = i;
			}

		}

		BufferedWriter bw = null;
		try {
			tags.add(tagSet.get(indexForMax));
			int lastIndex = indexForMax;

			bw = new BufferedWriter(new FileWriter("out.pos"));

			for (int i = n - 2; i >= 0; i--) {

				String tag = tagSet.get(bptr[lastIndex][i + 1]);
				tags.add(tag);
				lastIndex = bptr[lastIndex][i];
			}

			Collections.reverse(tags);

			for (int i = 0; i < n; i++) {
				bw.write(tags.get(i));
				bw.write(' ');
				bw.write(inputWords.get(i));
				bw.newLine();
			}

		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} finally {
			// Close the BufferedWriter
			try {
				if (bw != null) {
					bw.flush();
					bw.close();
				}
			} catch (IOException ex) {
				ex.printStackTrace();
			}
		}

		return tags;
	}
}
